# 进行期限套利操作

from trade.strategy import BaseStrategy
from trade.tools import SlidingData
from typing import Dict, List


class SpotFutureProfit(BaseStrategy):
    positions: dict[str, dict]

    def __init__(self, consider_futures: List[str]):
        self.positions = None
        self.consider_futures = consider_futures[:]
        self.spreads: Dict[str, SlidingData] = {
            x: SlidingData(10, [5]) for x in consider_futures
        }
        self.total_kinds = 5

    def apply(self):
        positions = self.broker.positions

        use_kinds = []
        for future in self.consider_futures:
            base = (
                self.broker.markets[future]["base"]
                + "/"
                + self.broker.markets[future]["quote"]
            )
            base_price = self.broker.get_data(base, "close", 0)
            ts = self.broker.get_data(future, "close", 0) - base_price
            self.spreads[future].push(ts)
            (m2, s2), (m1, s1) = self.spreads[future].get_value(5, slice(-2))
            if not (m1 or m2):
                continue

            tp = m1 / base_price

            if m1 > m2 and tp < 0:  # 低点回转
                if future in positions and positions[future]["side"] == "short":
                    size = positions[future]["size"]
                    self.broker.close_position(future, None, "GTC", "market")
                    self.broker.order_spot(base, "buy", size, None, "GTC", "market")
                if tp < -3e-3 and future not in positions:
                    use_kinds.append((future, abs(tp), "buy"))

            if m1 < m2 and tp > 0:
                if future in positions and positions[future]["side"] == "long":
                    size = positions[future]["size"]
                    self.broker.close_position(future, None, "GTC", "market")
                    self.broker.order_spot(base, "sell", size, None, "GTC", "market")
                if tp > 3e-3 and future not in positions:
                    use_kinds.append((future, abs(tp), "sell"))

        use_kinds = sorted(use_kinds, key=lambda x: x[1], reverse=True)
        need_new = self.total_kinds - len(self.positions)
        cur_asset = self.broker.get_account_value()
        if need_new > 0:
            for i in range(need_new):
                future, _, direction = use_kinds[i]
                base = (
                    self.broker.markets[future]["base"]
                    + "/"
                    + self.broker.markets[future]["quote"]
                )
                size = (
                    cur_asset
                    * 0.9
                    / self.total_kinds
                    / 2
                    / self.broker.get_data(future, "close", 0)
                )
                if direction == "buy":
                    self.broker.order_future(future, "buy", size, None, "GTC", "market")
                    self.broker.order_spot(base, "sell", size, None, "GTC", "market")
                else:
                    self.broker.order_future(
                        future, "sell", size, None, "GTC", "market"
                    )
                    self.broker.order_spot(base, "buy", size, None, "GTC", "market")
